from transformers import BertTokenizer, BertModel
from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader
import torch
import pandas as pd
from utils.remove_punc import pretreatment
from utils.config import ModelConfig
from model.train import train_model
from model.test import test_model
model_config = ModelConfig()
# data = pd.read_csv(model_config.train_path, encoding='utf-8')
# train_data = pretreatment(list(data['comment'].values))
train_df = pd.read_csv(model_config.train_path, sep='\t', header=0)
train_data = pretreatment(list(train_df['text_a'].values))

val_df = pd.read_csv(model_config.dev_path, sep='\t', header=0)
val_data = pretreatment(list(val_df['text_a'].values))

tokenizer = BertTokenizer.from_pretrained(model_config.bert_path)
train_ds_id = tokenizer(train_data,
                        padding=True,
                        truncation=True,
                        max_length=220,
                        return_tensors='pt')

val_ds_id = tokenizer(val_data,
                      padding=True,
                      truncation=True,
                      max_length=220,
                      return_tensors='pt')

X = train_ds_id['input_ids']
y = torch.from_numpy(train_df['label'].values).float()

X_test = val_ds_id['input_ids']
y_test = torch.from_numpy(val_df['label'].values).float()

X_train, _, y_train, _ = train_test_split(X,
                                          y,
                                          test_size=0.0005,
                                          shuffle=True,
                                          stratify=y,
                                          random_state=0)
print("---------",len(X_train))
X_valid, _, y_valid, _ = train_test_split(X_test,
                                          y_test,
                                          test_size=0.9,
                                          shuffle=True,
                                          stratify=y_test,
                                          random_state=0)
train_data = TensorDataset(X_train, y_train)
valid_data = TensorDataset(X_valid, y_valid)

train_loader = DataLoader(train_data,
                          shuffle=True,
                          batch_size=model_config.batch_size,
                          drop_last=True)
valid_loader = DataLoader(valid_data,
                          shuffle=True,
                          batch_size=model_config.batch_size,
                          drop_last=True)

if __name__ == "__main__":
    train_model(model_config, train_loader, valid_loader)
    # test_model(model_config, valid_loader)